package com.choudou5.solr.util;

import com.choudou5.solr.analyzer.bean.TermBean;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.analysis.tokenattributes.OffsetAttribute;
import org.apache.lucene.analysis.tokenattributes.TypeAttribute;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.wltea.analyzer.lucene.IKAnalyzer;

import java.io.IOException;
import java.io.StringReader;
import java.util.*;

/**
 * @Name：分词 工具类
 * @Author：choudou5
 * @Date：2018-07-24
 */
public class WordSegmUtil {

    private static Logger logger = LoggerFactory.getLogger(WordSegmUtil.class);

    private static final Analyzer SMART_ANZLYZER = new IKAnalyzer(true);

    public static void main(String[] args) {
        segmText("拼多多真的很多有用的词语");
    }

    public static String segmText(String text){
        return segmText(text, 8);
    }

    /**
     * 分割文本
     * @param text
     * @return
     */
    public static String segmText(String text, int appendScore){
        TokenStream ts = null;
        try {
            ts = SMART_ANZLYZER.tokenStream("text", new StringReader(text));
            OffsetAttribute offset = (OffsetAttribute)ts.addAttribute(OffsetAttribute.class);
            CharTermAttribute term = (CharTermAttribute)ts.addAttribute(CharTermAttribute.class);
            TypeAttribute type = (TypeAttribute)ts.addAttribute(TypeAttribute.class);
            ts.reset();
            StringBuilder expr = new StringBuilder(text.length()+8);
            int dex = 0;
            while (ts.incrementToken()) {
                if(dex > 0)
                    expr.append(" ");
                expr.append(term.toString());
                if(type.type().equals("CN_WORD") && appendScore > 0){
                    if(appendScore > 0){
                        expr.append("^"+appendScore);
                        appendScore -= 2;
                    }
                }
                System.out.println(offset.startOffset() + " - " + offset.endOffset() + " : " + term.toString() + " | " + type.type());
                dex++;
            }
            ts.end();
            System.out.println(expr.toString());
            return expr.toString();
        } catch (IOException e) {
            logger.error("segmText error", e);
            if (ts != null)
                try {
                    ts.close();
                } catch (IOException e2) {
                    e2.printStackTrace();
                }
        } finally  {
            if (ts != null)
                try {
                    ts.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
        }
        return text;
    }

    /**
     * 获取 词语 列表
     * @param text
     * @return
     */
    public static Set<String> getCNTermList(String text){
        TokenStream ts = null;
        try {
            ts = SMART_ANZLYZER.tokenStream("text", new StringReader(text));
            CharTermAttribute term = (CharTermAttribute)ts.addAttribute(CharTermAttribute.class);
            TypeAttribute type = (TypeAttribute)ts.addAttribute(TypeAttribute.class);
            ts.reset();
            Set<String> terms = new HashSet<>();
            while (ts.incrementToken()) {
                if(type.type().equals("CN_WORD")){
                    terms.add(term.toString());
                }
            }
            ts.end();
            return terms;
        } catch (IOException e) {
            logger.error("segmText getCNTermList", e);
            if (ts != null)
                try {
                    ts.close();
                } catch (IOException e2) {
                    e2.printStackTrace();
                }
        } finally  {
            if (ts != null)
                try {
                    ts.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
        }
        return Collections.EMPTY_SET;
    }

    /**
     * 获取 词语 列表
     * @param text
     * @return
     */
    public static Set<TermBean> getTermList(String text){
        TokenStream ts = null;
        try {
            ts = SMART_ANZLYZER.tokenStream("text", new StringReader(text));
            CharTermAttribute term = (CharTermAttribute)ts.addAttribute(CharTermAttribute.class);
            TypeAttribute type = (TypeAttribute)ts.addAttribute(TypeAttribute.class);
            ts.reset();
            Set<TermBean> terms = new HashSet<>();
            while (ts.incrementToken()) {
                terms.add(new TermBean(term.toString(), type.type()));
            }
            ts.end();
            return terms;
        } catch (IOException e) {
            logger.error("getTermList", e);
            if (ts != null)
                try {
                    ts.close();
                } catch (IOException e2) {
                    e2.printStackTrace();
                }
        } finally  {
            if (ts != null)
                try {
                    ts.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
        }
        return Collections.EMPTY_SET;
    }
}
